package io.focus.mijin.zuul.ratelimit.filters;

import com.netflix.zuul.context.RequestContext;
import io.focus.mijin.zuul.ratelimit.config.Rate;
import io.focus.mijin.zuul.ratelimit.config.RateLimitKeyGenerator;
import io.focus.mijin.zuul.ratelimit.config.RateLimiter;
import io.focus.mijin.zuul.ratelimit.config.properties.RateLimitProperties;
import io.focus.mijin.zuul.ratelimit.support.RateLimitExceededException;
import io.focus.mijin.zuul.ratelimit.support.RateLimitUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.netflix.zuul.filters.Route;
import org.springframework.cloud.netflix.zuul.filters.RouteLocator;
import org.springframework.http.HttpStatus;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.util.UrlPathHelper;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

import static io.focus.mijin.zuul.ratelimit.support.RateLimitConstants.*;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.FORM_BODY_WRAPPER_FILTER_ORDER;
import static org.springframework.cloud.netflix.zuul.filters.support.FilterConstants.PRE_TYPE;
import static org.springframework.web.context.request.RequestAttributes.SCOPE_REQUEST;

/**
 * 请求发起前调用的拦截器
 *
 * @author Wuaj
 */
public class RateLimitPreFilter extends AbstractRateLimitFilter {

    private static final Logger LOGGER = LoggerFactory.getLogger(RateLimitPreFilter.class);

    private final RateLimiter rateLimiter;

    private final RateLimitKeyGenerator rateLimitKeyGenerator;

    public RateLimitPreFilter(RateLimitProperties properties, RouteLocator routeLocator,
                              UrlPathHelper urlPathHelper, RateLimiter rateLimiter, RateLimitKeyGenerator rateLimitKeyGenerator,
                              RateLimitUtils rateLimitUtils) {
        super(properties, routeLocator, urlPathHelper, rateLimitKeyGenerator, rateLimitUtils);
        this.rateLimiter = rateLimiter;
        this.rateLimitKeyGenerator = rateLimitKeyGenerator;
    }

    @Override
    public String filterType() {
        return PRE_TYPE;
    }

    @Override
    public int filterOrder() {
        //return 1;
        return FORM_BODY_WRAPPER_FILTER_ORDER + 1;
    }

    @Override
    public Object run() {
        LOGGER.info("Pass Through RateLimitPreFilter Gateway Service Start");
        final RequestContext ctx = RequestContext.getCurrentContext();
        final HttpServletResponse response = ctx.getResponse();
        final HttpServletRequest request = ctx.getRequest();
        final Route route = route(request);

        policy(route, request).forEach(policy -> {
            LOGGER.info(" RateLimitPreFilter Check ");
            final String key = rateLimitKeyGenerator.key(request, route, policy);
            final Rate rate = rateLimiter.consume(policy, key, null);
            final String httpHeaderKey = key.replaceAll("[^A-Za-z0-9-.]", "_").replaceAll("__", "_");

            final Long limit = policy.getLimit();
            final Long remaining = rate.getRemaining();
            if (limit != null) {
                response.setHeader(HEADER_LIMIT + httpHeaderKey, String.valueOf(limit));
                response.setHeader(HEADER_REMAINING + httpHeaderKey, String.valueOf(Math.max(remaining, 0)));
            }

            final Long quota = policy.getQuota();
            final Long remainingQuota = rate.getRemainingQuota();
            if (quota != null) {
                RequestContextHolder.getRequestAttributes()
                        .setAttribute(REQUEST_START_TIME, System.currentTimeMillis(), SCOPE_REQUEST);
                response.setHeader(HEADER_QUOTA + httpHeaderKey, String.valueOf(quota));
                response.setHeader(HEADER_REMAINING_QUOTA + httpHeaderKey,
                        String.valueOf(MILLISECONDS.toSeconds(Math.max(remainingQuota, 0))));
            }

            response.setHeader(HEADER_RESET + httpHeaderKey, String.valueOf(rate.getReset()));

            boolean isOver = (limit != null && remaining < 0) || (quota != null && remainingQuota < 0);
            if (isOver) {
                ctx.setResponseStatusCode(HttpStatus.TOO_MANY_REQUESTS.value());
                ctx.put("rateLimitExceeded", "true");
                ctx.setSendZuulResponse(false);
                LOGGER.info(" 请求数量超过限制 ");
                throw new RateLimitExceededException();
            }
        });
        LOGGER.info("Pass Through RateLimitPreFilter Gateway Service End");
        return null;
    }
}
